package org.icatproject.ijp.scarf.ejb.session;

import java.io.BufferedWriter;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintStream;
import java.io.StringWriter;
import java.nio.file.FileSystems;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import javax.annotation.PostConstruct;
import javax.ejb.EJB;
import javax.ejb.Schedule;
import javax.ejb.Stateless;
import javax.json.Json;
import javax.json.JsonArray;
import javax.json.JsonArrayBuilder;
import javax.json.JsonObject;
import javax.json.JsonWriter;
import javax.persistence.EntityManager;
import javax.persistence.PersistenceContext;
import javax.xml.bind.JAXBContext;
import javax.xml.bind.Unmarshaller;

import org.icatproject.ICAT;
import org.icatproject.IcatException_Exception;
import org.icatproject.ijp.scarf.Bjobs;
import org.icatproject.ijp.scarf.Families;
import org.icatproject.ijp.scarf.Icat;
import org.icatproject.ijp.scarf.ejb.entity.Account;
import org.icatproject.ijp.scarf.ejb.entity.Job;
import org.icatproject.ijp.shared.AccountDTO;
import org.icatproject.ijp.shared.Constants;
import org.icatproject.ijp.shared.ForbiddenException;
import org.icatproject.ijp.shared.InternalException;
import org.icatproject.ijp.shared.ParameterException;
import org.icatproject.ijp.shared.PortalUtils.OutputType;
import org.icatproject.ijp.shared.SessionException;
import org.icatproject.utils.ShellCommand;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Session Bean implementation to manage job status
 */
@Stateless
public class JobManagementBean {

	private ICAT icat;

	private class LocalFamily {

		private int count;
		private Pattern re;

		public LocalFamily(int count, Pattern re) {
			this.count = count;
			this.re = re;
		}

	}

	@EJB
	private MachineEJB machineEJB;

	private String defaultFamily;
	private Map<String, LocalFamily> families = new HashMap<String, LocalFamily>();

	@PostConstruct
	void init() {
		try {
			icat = Icat.getIcat();

			Unmarshaller um = JAXBContext.newInstance(Families.class).createUnmarshaller();
			Families fams = (Families) um.unmarshal(new FileReader(Constants.CONFIG_SUBDIR
					+ "/families.xml"));
			PrintStream p = new PrintStream("/etc/puppet/modules/usergen/manifests/init.pp");
			p.print(fams.getPuppet());
			p.close();
			defaultFamily = fams.getDefault();
			for (Families.Family fam : fams.getFamily()) {
				LocalFamily lf = new LocalFamily(fam.getCount(), fam.getRE());
				families.put(fam.getName(), lf);
			}

			logger.debug("Initialised JobManagementBean");
		} catch (Exception e) {
			String msg = e.getClass().getName() + " reports " + e.getMessage();
			logger.error(msg);
			throw new RuntimeException(msg);
		}
	}

	private final static Logger logger = LoggerFactory.getLogger(JobManagementBean.class);
	private final static Random random = new Random();
	private final static String chars = "abcdefghijklmnpqrstuvwxyz";
	
	// TODO Replace following statics with configuration parameters, or separate user pool management
	
	private final static String IJP_USER_PREFIX = "ijp";
	private final static int IJP_USER_COUNT = 2;

	@PersistenceContext(unitName = "ijp")
	private EntityManager entityManager;

	public List<Job> getJobsForUser(String sessionId) throws SessionException {
		String username = getUserName(sessionId);
		return entityManager.createNamedQuery(Job.FIND_BY_USERNAME, Job.class)
				.setParameter("username", username).getResultList();
	}

	public String getJobOutput(String sessionId, String jobId, OutputType outputType)
			throws SessionException, ForbiddenException, InternalException {
		Job job = getJob(sessionId, jobId);
		String ext = "." + (outputType == OutputType.STANDARD_OUTPUT ? "log" : "err");
		Path path = FileSystems.getDefault().getPath(Constants.SCARF_POOL_BASE, job.getBatchUsername(),
				job.getId() + ext);
		boolean delete = false;
		if (!Files.exists(path)) {
			logger.debug("Getting intermediate output for " + jobId);
			
			// TODO Not sure what the equivalent of this is in Scarf; this is R92-specific. Perhaps use bpeek?
			
			ShellCommand sc = new ShellCommand("sudo", "-u", "batch", "ssh", job.getWorkerNode(),
					"sudo", "push_output", job.getBatchUsername(), path.toFile().getName());
			if (sc.isError()) {
				return "Temporary? problem getting output " + sc.getStderr();
			}
			path = FileSystems.getDefault().getPath(Constants.SCARF_POOL_BASE, job.getBatchUsername(),
					job.getId() + ext + "_tmp");
			delete = true;
		}
		if (Files.exists(path)) {
			logger.debug("Returning output for " + jobId);
			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			try {
				Files.copy(path, baos);
			} catch (IOException e) {
				throw new InternalException(e.getClass() + " reports " + e.getMessage());
			}
			if (delete) {
				try {
					Files.deleteIfExists(path);
				} catch (IOException e) {
					throw new InternalException("Unable to delete temporary file");
				}
			}
			return baos.toString();

		} else {
			throw new InternalException("No output file available at the moment");
		}
	}

	@Schedule(minute = "*/1", hour = "*")
	public void updateJobsFromBjobs() {
		try {

			// Need to run bjobs for *each* (active) user in the pool
			
			for( String poolUserId : getActivePoolUsers() ){
			
				ShellCommand sc = new ShellCommand("bjobs", "-aw", "-u", poolUserId );
				if (sc.isError()) {
					throw new InternalException("Unable to query jobs via bjobs " + sc.getStderr());
				}
				String bJobsOutput = sc.getStdout().trim();
				if (bJobsOutput.isEmpty()) {
					/* See if any jobs have completed without being noticed */
					for (Job job : entityManager.createNamedQuery(Job.FIND_INCOMPLETE, Job.class)
							.getResultList()) {
						logger.warn("Updating status of job '" + job.getId() + "' from '"
								+ job.getStatus() + "' to 'C' as not known to bjobs");
						job.setStatus("C");
					}
					return;
				}
	
				Bjobs bJobs = new Bjobs( bJobsOutput );
				for (Bjobs.Job bjob : bJobs.getJobs()) {
					String id = bjob.getJobId();
					String status = bjob.getStatus();
					String wn = bjob.getExecHost();
					String workerNode = wn != null ? wn.split("/")[0] : "";
					
					// No comment returned by bjobs, so ignore this field
					// String comment = bjob.getComment() == null ? "" : bjob.getComment();
	
					Job job = entityManager.find(Job.class, id);
					if (job != null) {/* Log updates on portal jobs */
						if (!job.getStatus().equals(bjob.getStatus())) {
							logger.debug("Updating status of job '" + id + "' from '" + job.getStatus()
									+ "' to '" + status + "'");
							job.setStatus(status);
						}
						if (!job.getWorkerNode().equals(workerNode)) {
							logger.debug("Updating worker node of job '" + id + "' from '"
									+ job.getWorkerNode() + "' to '" + workerNode + "'");
							job.setWorkerNode(workerNode);
						}
					}
				}
			} // for activePoolUsers
		} catch (Exception e) {
			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			e.printStackTrace(new PrintStream(baos));
			logger.error("Update of db jobs from bjobs failed. Class " + e.getClass() + " reports "
					+ e.getMessage() + baos.toString());
		}
	}

	public String submitBatch(String sessionId, String executable, List<String> parameters, String reqFamily)
			throws ParameterException, SessionException, InternalException {
		
		String username;
		
		// TODO TEMPORARY HACK: bypass sessionId so we can try this out without getting ICAT involved!
		
		if( ! "magic".equals(sessionId) ){
			String family = reqFamily == null ? defaultFamily : reqFamily;
			LocalFamily lf = families.get(family);
			if (lf == null) {
				throw new ParameterException("Requested family " + reqFamily + " not known");
			}
			username = getUserName(sessionId);
	
			if (lf.re != null && !lf.re.matcher(username).matches()) {
				throw new ParameterException(username + " is not allowed to use family " + family);
			}
			// TODO Adapt Family to be a source of the Scarf user pool IDs?
			
			// String owner = family + random.nextInt(lf.count);
		} else {
			
			// We don't really use the username at present, so set it to junk
			username = "magic";
		}
				
		String owner = assignScarfIdTo( username );
		String idFileName = getScarfIdFileNameFor( owner );
		
		String queueName = Constants.IJP_DEFAULT_QUEUE;
		
		// For now, use the executable as the job name.  May want to improve on this later.
		
		String jobName = executable;

		/*
		 * The batch script needs to be written to disk by the dmf user (running glassfish) before
		 * it can be submitted via the qsub command as a less privileged batch user. First generate
		 * a unique name for it.
		 */
		File batchScriptFile = null;
		do {
			char[] pw = new char[10];
			for (int i = 0; i < pw.length; i++) {
				pw[i] = chars.charAt(random.nextInt(chars.length()));
			}
			String batchScriptName = new String(pw) + ".sh";
			batchScriptFile = new File(Constants.DMF_WORKING_DIR_NAME, batchScriptName);
		} while (batchScriptFile.exists());

		createScript(batchScriptFile, jobName, executable, parameters, sessionId);
		
		ShellCommand sc = new ShellCommand("ssh", "-i", idFileName, owner + "@localhost", "bsub", "-J", jobName, "-o", "%J.log", "-e", "%J.err", "-q", queueName,
				batchScriptFile.getAbsolutePath());
		if (sc.isError()) {
			throw new InternalException("Unable to submit job via bsub " + sc.getStderr());
		}
		String jobId = getJobId(sc.getStdout());
		
		if( jobId == null ){
			throw new InternalException("Unable to get JobId for submitted job " + jobName);
		}

		// Query initial status of this job
		
		sc = new ShellCommand("ssh", "-i", idFileName, owner + "@localhost", "bjobs", "-aw", jobId);
		if (sc.isError()) {
			throw new InternalException("Unable to query just submitted job (id " + jobId
					+ ") via bjobs " + sc.getStderr());
		}
		
		Bjobs bjobs = new Bjobs( sc.getStdout() );
		Bjobs.Job bjob = bjobs.getJob(jobId);
		if( bjob == null ){
			throw new InternalException("Unable to find just submitted job (id " + jobId + ") in bjobs" );
		}
		
		Job job = new Job();
		job.setId(jobId);
		job.setStatus(bjob.getStatus());
		
		// Original processing: may want to reproduce here
		// String wn = xjob.getWorkerNode();
		// job.setWorkerNode(wn != null ? wn.split("/")[0] : "");

		job.setFromHost(bjob.getFromHost());
		job.setWorkerNode(bjob.getExecHost());
		job.setBatchUsername(bjob.getUserId());
		job.setSubmitDate(bjob.getSubmitTime());
		job.setBatchFilename(batchScriptFile.getName());
		entityManager.persist(job);

		return jobId;
	}

	private String getScarfIdFileNameFor(String owner) {
		return "$HOME/.ssh/id_rsa_" + owner;
	}

	/**
	 * Assign a Scarf userid to the given (ICAT) username.
	 * At present, the ID is chosen at random from a hard-wired pool;
	 * we ought to maintain a map from username to assigned poolid,
	 * so that subsequent job submissions from the same username
	 * will use the same Scarf userid.
	 * Once a job exists, we should always use its batchUsername
	 * to get the Scarf userid that submitted the job.
	 * 
	 * @param username
	 * @return Scarf userid
	 */
	private String assignScarfIdTo(String username) {
		// TODO Maintain map from username to Scarf pool user  
		// For now, owner (Scarf user) is hard-coded, chosen at random here.
		// assume ids have format prefix<n>, where 1 <= n <= 99
		
		return IJP_USER_PREFIX + String.format("%02d", random.nextInt(IJP_USER_COUNT)+1);
	}

	private List<String> getActivePoolUsers() {
		
		// TODO use proper pool management; hard-coded for now
		
		List<String> activeUsers = new ArrayList<String>();
		
		for( int i = 1; i <= IJP_USER_COUNT; i++ ){
			activeUsers.add( IJP_USER_PREFIX + String.format("%02d", i) );
		}
		return activeUsers;
	}

	private String getJobId(String stdout) throws InternalException {
		/* 
		 * bsub output format is "Job <jobId> is submitted to queue <queueName>.",
		 * (including the angle brackets).
		 */
		Pattern p = Pattern.compile("Job <(\\d+)> is submitted to queue <(\\w+)>.");
		Matcher m = p.matcher(stdout);
		if( ! m.matches() ){
			throw new InternalException("Unable to get JobId for submitted job: bsub output was: " + stdout);
		}
		return m.group(1);
	}

	private void createScript(File batchScriptFile, String jobName, String executable, List<String> parameters, 
			String sessionId) throws InternalException {

		List<String> finalParameters = new ArrayList<String>();
		finalParameters.addAll(parameters);

		BufferedWriter bw = null;
		try {
			bw = new BufferedWriter(new FileWriter(batchScriptFile));
			bw.write("#!/bin/sh");
			bw.newLine();
			bw.write("echo $(date) - " + jobName + " starting");
			bw.newLine();
			String line = executable + " "
					+ JobManagementBean.escaped(finalParameters);
			logger.debug("Exec line for " + jobName + ": " + line);
			bw.write(line);
			bw.newLine();
			bw.newLine();
		} catch (IOException e) {
			throw new InternalException("Exception creating batch script: " + e.getMessage());
		} finally {
			if (bw != null) {
				try {
					bw.close();
				} catch (IOException e) {
					// Ignore it
				}
			}
		}
		batchScriptFile.setExecutable(true);

	}

	private static String sq = "\"'\"";

	static String escaped(List<String> parameters) {
		StringBuilder sb = new StringBuilder();
		for (String parameter : parameters) {
			if (sb.length() != 0) {
				sb.append(" ");
			}
			int offset = 0;
			while (true) {
				int quote = parameter.indexOf('\'', offset);
				if (quote == offset) {
					sb.append(sq);
				} else if (quote > offset) {
					sb.append("'" + parameter.substring(offset, quote) + "'" + sq);
				} else if (offset != parameter.length()) {
					sb.append("'" + parameter.substring(offset) + "'");
					break;
				} else {
					break;
				}
				offset = quote + 1;
			}
		}
		return sb.toString();
	}

	public AccountDTO submitInteractive(String sessionId, String executable, List<String> parameters, String family)
			throws InternalException {
		
		// TODO Determine proper interactive job process for Platform LSF - existing approach will not work.
		
		Path p = null;
		try {
			p = Files.createTempFile(null, null);
		} catch (IOException e) {
			throw new InternalException("Unable to create a temporary file: " + e.getMessage());
		}
		File interactiveScriptFile = p.toFile();
		
		// createScript() allows jobName to be different from executable, but we use the same for both here
		
		createScript(interactiveScriptFile, executable, executable, parameters, sessionId);
		Account account = machineEJB.prepareMachine(sessionId, executable, parameters,
				interactiveScriptFile);
		return account.getDTO(machineEJB.getPoolPrefix());
	}

	private String getUserName(String sessionId) throws SessionException {
		try {
			return icat.getUserName(sessionId);
		} catch (IcatException_Exception e) {
			throw new SessionException("IcatException " + e.getFaultInfo().getType() + " "
					+ e.getMessage());
		}
	}

	public String listStatus(String sessionId) throws SessionException {
		String username = getUserName(sessionId);
		List<Job> jobs = entityManager.createNamedQuery(Job.FIND_BY_USERNAME, Job.class)
				.setParameter("username", username).getResultList();
		StringBuilder sb = new StringBuilder();
		for (Job job : jobs) {
			sb.append(job.getId() + ", " + job.getStatus() + "\n");
		}
		return sb.toString();

	}

	/**
	 * Get the status (as a JSON string) of the given job for the given session.
	 * 
	 * @param jobId
	 * @param sessionId
	 * @return Job details as JSON string
	 * @throws SessionException
	 * @throws ForbiddenException
	 */
	public String getStatus(String jobId, String sessionId) throws SessionException,
			ForbiddenException {
		Job job = getJob(sessionId, jobId);
		String jobString = null;
		
		if( job != null ){
			JsonObject jsonJob = job.getJson();
			StringWriter stWriter = new StringWriter();
			try (JsonWriter jsonWriter = Json.createWriter(stWriter)) {
			   jsonWriter.writeObject(jsonJob);
			}

			jobString = stWriter.toString();
		}
		return jobString;
	}
	
	/**
	 * Get the status of all jobs for the given session's user, as a JSON array-string.
	 * 
	 * @param sessionId
	 * @return Details of all (this session's user's) jobs as a JSON array-string
	 * @throws SessionException
	 * @throws ForbiddenException
	 */
	public String getStatus(String sessionId) throws SessionException, ForbiddenException {
		
		JsonArrayBuilder jsonArrayBuilder = Json.createArrayBuilder();
		for( Job job : getJobsForUser(sessionId) ){
			jsonArrayBuilder.add(job.getJson());
		}
		
		JsonArray jobsJson = jsonArrayBuilder.build();
		StringWriter stWriter = new StringWriter();
		try (JsonWriter jsonWriter = Json.createWriter(stWriter)) {
		   jsonWriter.writeArray(jobsJson);
		}

		return stWriter.toString();
	}

	private Job getJob(String sessionId, String jobId) throws SessionException, ForbiddenException {
		String username = getUserName(sessionId);
		Job job = entityManager.find(Job.class, jobId);
		if (job == null || !job.getUsername().equals(username)) {
			throw new ForbiddenException("Job does not belong to you");
		}
		return job;
	}

	public String delete(String sessionId, String jobId) throws SessionException,
			ForbiddenException, InternalException, ParameterException {
		Job job = getJob(sessionId, jobId);
		if( !Arrays.asList(Job.STATUS_COMPLETE_VALS).contains(job.getStatus()) ){
			throw new ParameterException(
					"Only completed jobs can be deleted - try cancelling first");
		}
		for (String oe : new String[] { "o", "e" }) {
			String ext = "." + oe + jobId.split("\\.")[0];
			Path path = FileSystems.getDefault().getPath("/home/batch/jobs",
					job.getBatchFilename() + ext);
			try {
				Files.deleteIfExists(path);
			} catch (IOException e) {
				throw new InternalException("Unable to delete " + path.toString());
			}
		}
		entityManager.remove(job);
		return "";
	}

	public String cancel(String sessionId, String jobId) throws SessionException,
			ForbiddenException, InternalException {
		Job job = getJob(sessionId, jobId);
		String owner = job.getBatchUsername();
		String idFileName = getScarfIdFileNameFor( owner );
		ShellCommand sc = new ShellCommand("ssh", "-i", idFileName, owner + "@localhost", "bkill", job.getId());
		if (sc.isError()) {
			throw new InternalException("Unable to cancel job " + sc.getStderr());
		}
		return "";
	}

}
